//============================================================================
// Name        : Coevolution 2
// Author      : Alex Seeholzer
// Description :
// Adaptation of Benedikt Obermayers original code to two rival populations of an immune and a viral population
// This also assumes a functional and a structural divison of the populations.
//
// Usage:		 ./QuasiSpecies .0003 .0001 5. 5. .3 15. 25
//============================================================================

#include <stdio.h>
#include <math.h>
#include <typeinfo>
#include <gsl/gsl_math.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_interp.h>

#define TYPE unsigned int
#define L 32
#define ONE 1
#define RRMASK 0x000000FF
#define MAXINT 0xFFFFFFFF
#define LRR 8
#define VIR 0
#define IS 1
#define DELTA 1.
#define N 10000
#define TMAX 1000

#define NM 2 /* population number, 0=virus 1=immune */
const TYPE masters[] = {0x88888800, 0xFFFFFF00}; /* the individual master sequences for structure */

double alpha[2];
double gam[2];

inline TYPE mutate (TYPE x, double mu);
inline int hd (TYPE x, TYPE y); // hamming dist
inline void outbinary(TYPE num);
void reproduce (TYPE pop[NM][N], double* mu, unsigned int cdis[NM][(ONE << LRR)]);
void sample(TYPE* pop,TYPE* popnew,double* mut,double* sel,double mu);
TYPE compute_consensus (TYPE *pop);

// computes for a pop the distribution of clones (stored in cdis)
void compute_cdis (TYPE *pop, unsigned int *cdis, TYPE master);

// Random number generator
gsl_rng *r;
gsl_interp_accel *acc;
const gsl_rng_type *T;

int main(int argc, char **argv) {

	// Set up output file
	FILE * outfile;
	outfile = fopen("out.dat","w");

	int i,j,k,t,s;
	TYPE pop[NM][N]; // array of population
	double mu[NM]; // mutation rate
	TYPE smax[NM]; // maximal populated sequences for each population
	unsigned int cdis[NM][(ONE << LRR)]; // for each population, distribution of functional region

	// random number generators
	gsl_rng_env_setup();
	acc = gsl_interp_accel_alloc();
	T = gsl_rng_gfsr4;
	r = gsl_rng_alloc(T);

	// read mu, s = seed
	if (argc > 7) {
		mu[VIR]=atof(argv[1]);
		mu[IS]=atof(argv[2]);
		alpha[VIR]=atof(argv[3]);
		alpha[IS]=atof(argv[4]);
		gam[VIR]=atof(argv[5]);
		gam[IS]=atof(argv[6]);
		s=atoi(argv[7]);
		fprintf(stderr, "mu_v=% .3e, mu_is=% .3e, alpha_v=% .3e, alpha_is=% .3e, gamma_v=% .3e, gamma_is=% .3e, s=%i\n", mu[VIR], mu[IS], alpha[VIR], alpha[IS], gam[VIR], gam[IS], s);
	} else {
		fprintf(stderr, "Usage:\n QuasiSpecies mu_v mu_is alpha_vir alpha_is gamma_vir gamma_is seed\n mu_v: mutation rate for virus population\nmu_is: mutation rate for immune system population\n");
		exit(1);
	}

	gsl_rng_set (r, s); // initialize random num generator with seed s

	// population initialization, set as master sequence initially
	for (j = 0; j < NM; j++)
		for (i = 0; i < N; i++)
			pop[j][i] = (masters[j] & ~RRMASK)
					| gsl_rng_uniform_int(r, ONE << LRR);

	// compute frequency of master sequences for initial reproduce call
	for (j = 0; j < NM; j++)
		compute_cdis(pop[j], cdis[j], masters[j]);

	// main loop - reproduce, calc distribution and output
	for (t = 0; t < TMAX; t++) {
		reproduce(pop, mu, cdis);
		for (j = 0; j < NM; j++)
			compute_cdis(pop[j], cdis[j], masters[j]);

		for (k = 0; k < (ONE << LRR); k++) {
			fprintf(outfile,"%i %i ", t, k);
			for (j = 0; j < NM; j++)
				fprintf(outfile,"%i ", cdis[j][k]);
			fprintf(outfile,"\n");
		}
		fprintf(outfile,"\n");
		if (t%100 == 0)
			printf("time: %i\n",t);
	}


	for (j = 0; j < NM; j++) {
		smax[j] = 0;
		for (k = 1; k < (ONE << LRR); k++)
			if (cdis[j][k] > cdis[j][smax[j]])
				smax[j] = k;

		fprintf(stderr, "smax[%i]= %3i (%i), neighbors: ", j, smax[j],
				cdis[j][smax[j]]);
		for (k = 0; k < LRR; k++)
			fprintf(stderr, "%i ", smax[j] ^ (ONE << k));
		fprintf(stderr, "\n");
	}

	fclose(outfile);
	return 0;
}

void compute_cdis (TYPE *pop, unsigned int *cdis, TYPE master) {
  for (int k=0; k<(ONE << LRR); k++)
    cdis[k]=0;

  for (int i=0; i<N; i++)
    if ((pop[i] & ~RRMASK)==(master & ~RRMASK))
      cdis[(unsigned int)(pop[i] & RRMASK)]++;

  return;
}

void reproduce(TYPE pop[NM][N], double* mu, unsigned int cdis[NM][(ONE << LRR)]) {
	int i, j;
	TYPE popnew[NM][N];
	double sel[N];
	double mut[N];
	double delta;

	// reproduce virus population
	for (i = 0; i < N; i++) {
		// calculate percentage of genome pop[VIR][i] in pop[IS]
		delta = cdis[IS][(pop[VIR][i]) & RRMASK] / (double) N;
		mut[i] = 1.;
		sel[i] = (1/DELTA)-(gam[VIR]*delta);
		if ((pop[VIR][i] & ~RRMASK) == (masters[VIR] & ~RRMASK)) // if we have master sequence
			mut[i] = alpha[VIR];//-(alpha[VIR]*delta);
	}
	sample(pop[VIR],popnew[VIR],mut,sel,mu[VIR]);

	// reproduce immune system population
	for (i = 0; i < N; i++) {
		// calculate percentage of genome pop[IS][i] in pop[VIR]
		delta = cdis[VIR][(pop[IS][i]) & RRMASK] / (double) N;
		mut[i] = 1.;
		sel[i] = (1/DELTA);
		if ((pop[IS][i] & ~RRMASK) == (masters[IS] & ~RRMASK)) // if we have master sequence
			mut[i] = alpha[IS]+(gam[IS]*delta);
	}
	sample(pop[IS],popnew[IS],mut,sel,mu[IS]);

	for (j = 0; j < NM; j++)
		for (i = 0; i < N; i++)
			pop[j][i] = popnew[j][i];
	return;
}

void sample(TYPE* pop, TYPE* popnew, double* mut,double* sel,double mu) {
	double ptot;
	double select[2*N];
	int i,index;
	// sum of probabilities for normalization
	ptot = 0;
	for (i = 0; i < N; i++)
		ptot += sel[i] + mut[i];

	// evaluate probabilities to be selected
	// (even index) and to be selected+mutated
	// (odd indices)
	select[0] = sel[0] / ptot;
	select[1] = mut[0] / ptot + select[0];
	for (i = 1; i < N; i++) {
		select[2* i ] = sel[i] / ptot + select[2* i - 1];
		select[2* i + 1] = mut[i] / ptot + select[2* i ];
	}
	//get N random numbers in [0,1)
	for (i = 0;	i < N;i++) {
		index = gsl_interp_bsearch(select, gsl_rng_uniform(r), 0, 2*N - 2) + 1; //look up entry in select array

		// if index is odd, mutate this sequence with probability mu, otherwise copy
		popnew[i] = (index % 2 == 1 ? mutate(pop[index / 2], mu) : pop[index / 2]);
	}
	return;
}

inline TYPE mutate (TYPE x, double mu) {
  int i;
  TYPE out=x;
  for (i=0; i<L; i++)
    if (gsl_rng_uniform(r) < mu)
      out^=(ONE << i);
  return out;
}

inline int hd (TYPE x, TYPE y) {
  int dist=0;
  TYPE val=x ^ y;
  while (val) {
    ++dist;
    val &= val-1;
  }
  return dist;
}

inline void outbinary(TYPE num) {
	TYPE tmp = num;
	for(int i=0;i<L;i++) {
		int powr = int(pow(2,(L-i-1)));
		int res = tmp/powr;
		fprintf(stderr,"%i",res);
		tmp-=res*powr;
	}
	fprintf(stderr,"\n");
}
